#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PDF OCR文字识别工具 (使用PaddleOCR)
递归遍历指定文件夹，使用OCR识别所有PDF文件中的文字内容并保存到txt文件
"""

import os
import sys
from pathlib import Path
from datetime import datetime
import tempfile

try:
    from pdf2image import convert_from_path
    HAS_PDF2IMAGE = True
except ImportError:
    print("错误：未安装pdf2image！正在尝试安装...")
    os.system("pip install pdf2image")
    try:
        from pdf2image import convert_from_path
        HAS_PDF2IMAGE = True
    except:
        print("请手动安装: pip install pdf2image")
        HAS_PDF2IMAGE = False

try:
    from paddleocr import PaddleOCR
    HAS_PADDLEOCR = True
except ImportError:
    print("错误：未安装PaddleOCR！正在尝试安装...")
    os.system("pip install paddleocr")
    try:
        from paddleocr import PaddleOCR
        HAS_PADDLEOCR = True
    except:
        print("请手动安装: pip install paddleocr")
        HAS_PADDLEOCR = False

if not HAS_PDF2IMAGE or not HAS_PADDLEOCR:
    print("\n依赖库安装失败，请手动运行:")
    print("pip install pdf2image paddleocr")
    sys.exit(1)


class PDFOCRExtractor:
    """PDF OCR文字提取器"""
    
    def __init__(self):
        """初始化OCR引擎"""
        print("正在初始化PaddleOCR引擎...")
        # 使用中文OCR模型，关闭角度分类器以提高速度
        self.ocr = PaddleOCR(
            use_angle_cls=True,  # 使用角度分类器
            lang="ch",  # 中文
            show_log=False,  # 不显示详细日志
            use_gpu=False  # 不使用GPU，如果有GPU可以改为True
        )
        print("OCR引擎初始化完成！\n")
    
    def pdf_to_images(self, pdf_path):
        """将PDF转换为图片列表"""
        try:
            # 转换PDF为图片，dpi设置为200以获得较好的识别效果
            images = convert_from_path(
                pdf_path,
                dpi=200,  # 分辨率，越高识别越准确但速度越慢
                fmt='jpeg'
            )
            return images
        except Exception as e:
            print(f"  ⚠️  转换PDF到图片失败: {str(e)}")
            return []
    
    def ocr_image(self, image):
        """对单张图片进行OCR识别"""
        try:
            # 将PIL Image转换为numpy数组
            import numpy as np
            image_array = np.array(image)
            
            # PaddleOCR识别
            result = self.ocr.ocr(image_array, cls=True)
            
            if result is None or len(result) == 0:
                return ""
            
            # 提取文字内容
            text_lines = []
            for line in result:
                if line:
                    for item in line:
                        if len(item) >= 2:
                            text = item[1][0]  # 获取识别的文字
                            text_lines.append(text)
            
            return '\n'.join(text_lines)
        except Exception as e:
            print(f"  ⚠️  OCR识别失败: {str(e)}")
            return ""
    
    def extract_text_from_pdf(self, pdf_path):
        """从PDF文件中提取文字（使用OCR）"""
        print(f"  → 正在转换PDF为图片...")
        images = self.pdf_to_images(str(pdf_path))
        
        if not images:
            return "错误：无法转换PDF为图片"
        
        print(f"  → 共 {len(images)} 页，开始OCR识别...")
        
        all_text = []
        for page_num, image in enumerate(images, 1):
            print(f"    - 识别第 {page_num}/{len(images)} 页...", end='', flush=True)
            
            # OCR识别
            text = self.ocr_image(image)
            
            if text.strip():
                all_text.append(f"\n{'='*60}\n")
                all_text.append(f"第 {page_num} 页\n")
                all_text.append(f"{'='*60}\n")
                all_text.append(text)
                print(" ✓")
            else:
                all_text.append(f"\n第 {page_num} 页（未识别到文字）\n")
                print(" (空)")
        
        return ''.join(all_text) if all_text else "未识别到任何文字内容"


def find_all_pdf_files(root_folder):
    """递归查找文件夹中的所有PDF文件"""
    pdf_files = []
    root_path = Path(root_folder)
    
    if not root_path.exists():
        print(f"错误：文件夹不存在 - {root_folder}")
        return pdf_files
    
    if not root_path.is_dir():
        print(f"错误：路径不是文件夹 - {root_folder}")
        return pdf_files
    
    # 递归查找所有PDF文件
    for pdf_file in root_path.rglob("*.pdf"):
        if pdf_file.is_file():
            pdf_files.append(pdf_file)
    
    # 同时查找PDF大写扩展名
    for pdf_file in root_path.rglob("*.PDF"):
        if pdf_file.is_file() and pdf_file not in pdf_files:
            pdf_files.append(pdf_file)
    
    return sorted(pdf_files)


def process_pdfs(folder_path, output_file=None):
    """处理文件夹中的所有PDF文件"""
    print(f"开始扫描文件夹: {folder_path}")
    print("-" * 60)
    
    # 查找所有PDF文件
    pdf_files = find_all_pdf_files(folder_path)
    
    if not pdf_files:
        print("未找到任何PDF文件！")
        return
    
    print(f"找到 {len(pdf_files)} 个PDF文件\n")
    
    # 如果没有指定输出文件，使用默认名称
    if output_file is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_file = f"pdf_ocr_extracted_text_{timestamp}.txt"
    
    # 初始化OCR提取器
    extractor = PDFOCRExtractor()
    
    # 提取文字并写入文件
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(f"PDF OCR文字提取结果\n")
        f.write(f"提取时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"源文件夹: {folder_path}\n")
        f.write(f"PDF文件数量: {len(pdf_files)}\n")
        f.write(f"OCR引擎: PaddleOCR (中文)\n")
        f.write(f"\n{'#'*60}\n\n")
        
        for idx, pdf_path in enumerate(pdf_files, 1):
            print(f"\n{'='*60}")
            print(f"[{idx}/{len(pdf_files)}] 正在处理: {pdf_path.name}")
            print(f"{'='*60}")
            
            # 写入文件信息
            f.write(f"\n{'#'*60}\n")
            f.write(f"文件 {idx}: {pdf_path.name}\n")
            f.write(f"完整路径: {pdf_path}\n")
            f.write(f"{'#'*60}\n")
            
            # 提取文字
            try:
                text = extractor.extract_text_from_pdf(pdf_path)
                
                if text.strip():
                    f.write(text)
                    f.write("\n\n")
                else:
                    f.write("（该PDF文件未识别到文字内容）\n\n")
                
                # 实时刷新文件，防止程序中断导致数据丢失
                f.flush()
                
            except Exception as e:
                error_msg = f"错误：处理PDF时发生异常 - {str(e)}\n\n"
                print(f"  ❌ {error_msg}")
                f.write(error_msg)
                f.flush()
    
    print("\n" + "="*60)
    print(f"✅ 完成！所有文字已保存到: {output_file}")
    print(f"📍 输出文件完整路径: {os.path.abspath(output_file)}")
    print("="*60)


def main():
    """主函数"""
    print("="*60)
    print("PDF OCR文字识别工具 (PaddleOCR)")
    print("="*60)
    print()
    
    # 获取文件夹路径
    if len(sys.argv) > 1:
        folder_path = sys.argv[1]
    else:
        folder_path = input("请输入要扫描的文件夹路径: ").strip()
    
    # 去除引号（如果用户输入时带了引号）
    folder_path = folder_path.strip('"').strip("'")
    
    # 获取输出文件名（可选）
    output_file = None
    if len(sys.argv) > 2:
        output_file = sys.argv[2]
    
    print()
    
    # 处理PDF文件
    try:
        process_pdfs(folder_path, output_file)
    except KeyboardInterrupt:
        print("\n\n⚠️  程序被用户中断")
    except Exception as e:
        print(f"\n❌ 发生错误: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()

